from __future__ import absolute_import

import os
import sys

# import numpy as np
# import scipy.misc
# import tensorflow as tf
from tensorboardX import SummaryWriter

try:
    from StringIO import StringIO  # Python 2.7
except ImportError:
    from io import BytesIO  # Python 3.x

from .distributed_ops import is_main_process
from .osutils import mkdir_if_missing

from config import get_args

global_args = get_args(sys.argv[1:])

if global_args.run_on_remote:
    import moxing as mox

    mox.file.shift("os", "mox")


class Logger(object):
    def __init__(self, fpath=None):
        self.console = sys.stdout
        self.file = None
        if fpath is not None:
            if global_args.run_on_remote:
                dir_name = os.path.dirname(fpath)
                if not mox.file.exists(dir_name):
                    mox.file.make_dirs(dir_name)
                    print('=> making dir ', dir_name)
                self.file = mox.file.File(fpath, 'w')
                # self.file = open(fpath, 'w')
            else:
                mkdir_if_missing(os.path.dirname(fpath))
                self.file = open(fpath, 'w')

    def __del__(self):
        self.close()

    def __enter__(self):
        pass

    def __exit__(self, *args):
        self.close()

    def write(self, msg):
        self.console.write(msg)
        if self.file is not None:
            self.file.write(msg)

    def flush(self):
        self.console.flush()
        if self.file is not None:
            self.file.flush()
            os.fsync(self.file.fileno())

    def close(self):
        self.console.close()
        if self.file is not None:
            self.file.close()


class TBLogger:
    def __init__(self, log_dir="logs"):
        if is_main_process():
            os.makedirs(log_dir, exist_ok=True)
            self.writer = SummaryWriter(logdir=log_dir)

    def add_scalar(self, tag, value, step):
        if is_main_process():
            self.writer.add_scalar(tag, value, step)
            self.writer.flush()

    def close(self):
        if is_main_process():
            self.writer.close()

# class TFLogger(object):
#     def __init__(self, log_dir=None):
#         """Create a summary writer logging to log_dir."""
#         if log_dir is not None:
#             mkdir_if_missing(log_dir)
#         # self.writer = tf.summary.FileWriter(log_dir)
#         self.writer = tf.summary.create_file_writer(log_dir)
#
#     def scalar_summary(self, tag, value, step):
#         """Log a scalar variable."""
#         # summary = tf.Summary(value=[tf.Summary.Value(tag=tag, simple_value=value)])
#         # self.writer.add_summary(summary, step)
#         # self.writer.flush()
#         with self.writer.as_default():
#             tf.summary.scalar(tag, value, step=step)
#             self.writer.flush()
#
#     def image_summary(self, tag, images, step):
#         """Log a list of images."""
#
#         img_summaries = []
#         for i, img in enumerate(images):
#             # Write the image to a string
#             try:
#                 s = StringIO()
#             except:
#                 s = BytesIO()
#             scipy.misc.toimage(img).save(s, format="png")
#
#             # Create an Image object
#             img_sum = tf.Summary.Image(encoded_image_string=s.getvalue(),
#                                        height=img.shape[0],
#                                        width=img.shape[1])
#             # Create a Summary value
#             img_summaries.append(tf.Summary.Value(tag='%s/%d' % (tag, i), image=img_sum))
#
#         # Create and write Summary
#         summary = tf.Summary(value=img_summaries)
#         self.writer.add_summary(summary, step)
#         self.writer.flush()
#
#     def histo_summary(self, tag, values, step, bins=1000):
#         """Log a histogram of the tensor of values."""
#
#         # Create a histogram using numpy
#         counts, bin_edges = np.histogram(values, bins=bins)
#
#         # Fill the fields of the histogram proto
#         hist = tf.HistogramProto()
#         hist.min = float(np.min(values))
#         hist.max = float(np.max(values))
#         hist.num = int(np.prod(values.shape))
#         hist.sum = float(np.sum(values))
#         hist.sum_squares = float(np.sum(values ** 2))
#
#         # Drop the start of the first bin
#         bin_edges = bin_edges[1:]
#
#         # Add bin edges and counts
#         for edge in bin_edges:
#             hist.bucket_limit.append(edge)
#         for c in counts:
#             hist.bucket.append(c)
#
#         # Create and write Summary
#         summary = tf.Summary(value=[tf.Summary.Value(tag=tag, histo=hist)])
#         self.writer.add_summary(summary, step)
#         self.writer.flush()
#
#     def close(self):
#         self.writer.close()
